#!/usr/bin/python3
import os
import socket
import struct
import threading
import time
import argparse


is_connected = False


def create_msg(msg):
    bytes_msg = msg.encode("utf-8")
    bodylen = socket.htonl(len(bytes_msg))
    id = socket.htonl(1000)
    return struct.pack("II", id, bodylen) + bytes_msg


def rev_msg(s: socket):
    while True:
        if protocol == "tcp":
            header = s.recv(8)
            if not header:
                print("server disconnected")
                is_connected = False
                return
            else:
                id, bodylen = struct.unpack("II", header)
                reallen = socket.ntohl(bodylen)
                result = s.recv(reallen)
                print("recv from server :", result.decode("utf-8"))
        else:
            rspbuffer = s.recv(65536)
            print("recv len:", len(rspbuffer))
            id, bodylen = struct.unpack("II", rspbuffer[0:8])
            print("recv from server :", rspbuffer[8:].decode("utf-8"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="xncc echo client test")
    parser.add_argument("-s", type=str, required=True, help="echo server address")
    parser.add_argument("-p", type=str, required=True, help="protocol udp/tcp")
    args = parser.parse_args()
    server_addr = args.s
    server_ip = server_addr.split(":")[0]
    server_port = int(server_addr.split(":")[1])
    print("start connect success, ", server_ip, ":", server_port)
    protocol = args.p
    while True:
        if protocol == "udp":
            proto = socket.SOCK_DGRAM
        elif protocol == "tcp":
            proto = socket.SOCK_STREAM
        s = socket.socket(socket.AF_INET, proto)
        err = s.connect_ex((server_ip, server_port))
        if err != 0:
            print("connect ", server_ip, ":", server_port, " failed ", os.strerror(err))
            time.sleep(2)
        else:
            print("connect success, ", server_ip, ":", server_port)
            is_connected = True
            break
    recv_thread = threading.Thread(target=rev_msg, args=(s,))
    recv_thread.start()

    while True:
        if not is_connected:
            break
        msg = input()
        s.sendall(create_msg(msg))
